module BASICSemantic

open BASICAST
open BASICUtils
open System.Collections.Generic
 

let toFloat x line = 
    try
        float x
    with
    | e -> 
        failwithf "Error in line: %d. Number too large: %s" line x

let joinMap (p:Map<'a,'b>) (q:Map<'a,'b>) = 
    Map(Seq.concat [ (Map.toSeq p) ; (Map.toSeq q) ])

let joinSet a b =
    a+b
 
let rec extractfuncFromExpr (expr : Expr) func joinFun =
    match expr with
        | IdentifierList(_,_,exp) -> joinFun (extractfuncFromExpr exp func joinFun) (func expr)
        | IdentifierTable(_,_,exp1, exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun) (joinFun (extractfuncFromExpr exp2 func joinFun)  (func expr))
        | Minus(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Plus(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Times(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Div(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Pow(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | MinEq(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | MaxEq(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | NotEq(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Min(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Max(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | Equal(_,exp1,exp2) -> joinFun (extractfuncFromExpr exp1 func joinFun)  (extractfuncFromExpr exp2 func joinFun)
        | UsrFun(_,_,exp) -> joinFun (extractfuncFromExpr exp func joinFun) (func expr)
        | StdFun(_,_,exp) -> joinFun (extractfuncFromExpr exp func joinFun) (func expr)
        | _ -> func(expr)

let extractScalarsFromExpr (expr : Expr) = 
    match expr with
        | Identifier(_, id) -> Set.empty.Add(id)
        | _ -> Set.empty


let extractArrayFromExpr (expr : Expr) = 
    match expr with
        | IdentifierList(_,id,_) -> Map.empty.Add(string id, 11)
        | _ -> Map.empty




let extractMatrixFromExpr (expr : Expr) = 
    match expr with
        | IdentifierTable(_,id,_,_) -> Map.empty.Add(string id, (11,11))
        | _ -> Map.empty


let extractFunctionFromExpr (expr : Expr) = 
    match expr with
        | UsrFun(num,id,_) -> Map.empty.Add(string id, num)
        | _ -> Map.empty


let extractScalarFromPrint scalarSet tupla =
    let t = match tupla with
            | (_,Some(exp),_) -> extractfuncFromExpr exp extractScalarsFromExpr joinSet
            | _ -> Set.empty
    scalarSet + t
    
let extractArrayFromPrint arrayMap tupla =
    let t = match tupla with
            | (_,Some(exp),_) -> extractfuncFromExpr exp extractArrayFromExpr joinMap
            | _ -> Map.empty
    joinMap arrayMap t



let extractFunctionFromPrint arrayMap tupla =
    let t = match tupla with
            | (_,Some(exp),_) -> extractfuncFromExpr exp extractFunctionFromExpr joinMap
            | _ -> Map.empty
    joinMap arrayMap t



let extractMatrixFromPrint matrixMap tupla =
    let t = match tupla with
            | (_,Some(exp),_) -> extractfuncFromExpr exp extractMatrixFromExpr joinMap
            | _ -> Map.empty
    joinMap matrixMap t

let extractScalarFromRead scalarSet tupla =
    let t = match tupla with
            | (_,Some(exp1),_) -> extractfuncFromExpr exp1 extractScalarsFromExpr joinSet
            | _ -> Set.empty
    let t2 = match tupla with
             | (_,_,Some(exp2)) -> extractfuncFromExpr exp2 extractScalarsFromExpr joinSet
             | _ -> Set.empty
    let t3 = match tupla with
             | (id, None, None) -> Set.empty.Add(id)
             | _ -> Set.empty
    scalarSet + (t + t2 + t3)


let extractArrayFromRead arrayMap tupla =
    let t = match tupla with
            | (_,Some(exp1),_) -> extractfuncFromExpr exp1 extractArrayFromExpr joinMap
            | _ -> Map.empty
    let t2 = match tupla with
             | (_,_,Some(exp2)) -> extractfuncFromExpr exp2 extractArrayFromExpr joinMap
             | _ -> Map.empty
    let t3 = match tupla with
             | (id, Some(_), None) -> Map.empty.Add(id, 11) 
             | _ -> Map.empty
    joinMap (joinMap arrayMap (joinMap t  t2)) t3



let extractFunctionFromRead arrayMap tupla =
    let t = match tupla with
            | (_,Some(exp1),_) -> extractfuncFromExpr exp1 extractFunctionFromExpr joinMap
            | _ -> Map.empty
    let t2 = match tupla with
             | (_,_,Some(exp2)) -> extractfuncFromExpr exp2 extractFunctionFromExpr joinMap
             | _ -> Map.empty
    joinMap arrayMap (joinMap t  t2)




let extractMatrixFromRead matrixMap tupla =
    let t = match tupla with
            | (_,Some(exp1),_) -> extractfuncFromExpr exp1 extractMatrixFromExpr joinMap
            | _ -> Map.empty
    let t2 = match tupla with
             | (_,_,Some(exp2)) -> extractfuncFromExpr exp2 extractMatrixFromExpr joinMap
             | _ -> Map.empty
    let t3 = match tupla with
             | (id, Some(_), Some(_)) -> Map.empty.Add(id, (11,11)) 
             | _ -> Map.empty
    joinMap (joinMap matrixMap (joinMap t  t2)) t3


let mutable dimList = Map.empty
let extractArrayFromDim arrayMap tupla =
    let t = match tupla with
            | (var,numAsStr,None) -> Map.empty.Add(var, int(numAsStr))
            | _ -> Map.empty
    dimList <- joinMap dimList t
    Map.empty


let mutable dimMatrix = Map.empty
let extractMatrixFromDim matrixMap tupla =
    let t = match tupla with
            | (var,numAsStr,Some(num2AsString)) -> Map.empty.Add(var, (int(numAsStr),int(num2AsString)))
            | _ -> Map.empty
    dimMatrix <- joinMap dimMatrix t
    Map.empty

let createScalarTable (scalarSet : Set<string>) stmt =
    let t = match stmt with
            | LetVar(_,_, id, exp) -> (extractfuncFromExpr exp extractScalarsFromExpr joinSet) + Set.empty.Add(id)
            | LetList(_,_,_, exp1, exp2) -> (extractfuncFromExpr exp1 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp2 extractScalarsFromExpr joinSet)
            | LetTable(_,_,_,exp1, exp2, exp3) -> (extractfuncFromExpr exp1 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp2 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp3 extractScalarsFromExpr joinSet)
            | If(_,_,exp1,_,exp2,_) -> (extractfuncFromExpr exp1 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp2 extractScalarsFromExpr joinSet)
            | Def(_,_,_,id,exp) -> (extractfuncFromExpr exp extractScalarsFromExpr joinSet) + Set.empty.Add(id)
            | Print(_,_, list) -> List.fold extractScalarFromPrint Set.empty list
            | Read(_,_,list) -> List.fold extractScalarFromRead Set.empty list
            | For(_,_,id,exp1,exp2, Some(exp3)) -> (extractfuncFromExpr exp1 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp2 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp3 extractScalarsFromExpr joinSet) + Set.empty.Add(id)
            | For(_,_,id,exp1,exp2, None) -> (extractfuncFromExpr exp1 extractScalarsFromExpr joinSet) + (extractfuncFromExpr exp2 extractScalarsFromExpr joinSet) + Set.empty.Add(id)
            | _ -> Set.empty
    scalarSet + t
    
    

let createArrayTable (arrayMap) stmt =
    let t = match stmt with
            | Dim (_,_,list) -> List.fold extractArrayFromDim Map.empty list
            | LetVar(_,_, _, exp) -> (extractfuncFromExpr exp extractArrayFromExpr joinMap)
            | LetList(_,_,id, exp1, exp2) -> joinMap (joinMap (extractfuncFromExpr exp1 extractArrayFromExpr joinMap) (extractfuncFromExpr exp2 extractArrayFromExpr joinMap)) (Map.empty.Add(string id, 11))
            | LetTable(_,_,_,exp1, exp2, exp3) -> joinMap (extractfuncFromExpr exp1 extractArrayFromExpr joinMap) (joinMap (extractfuncFromExpr exp2 extractArrayFromExpr joinMap) (extractfuncFromExpr exp3 extractArrayFromExpr joinMap))
            | If(_,_,exp1,_,exp2,_) -> joinMap (extractfuncFromExpr exp1 extractArrayFromExpr joinMap)  (extractfuncFromExpr exp2 extractArrayFromExpr joinMap)
            | Print(_,_, list) -> List.fold extractArrayFromPrint Map.empty list
            | Read(_,_,list) -> List.fold extractArrayFromRead Map.empty list
            | For(_,_,_,exp1,exp2, Some(exp3)) -> joinMap (extractfuncFromExpr exp1 extractArrayFromExpr joinMap) (joinMap (extractfuncFromExpr exp2 extractArrayFromExpr joinMap) (extractfuncFromExpr exp3 extractArrayFromExpr joinMap))
            | For(_,_,_,exp1,exp2, None) -> joinMap (extractfuncFromExpr exp1 extractArrayFromExpr joinMap)  (extractfuncFromExpr exp2 extractArrayFromExpr joinMap)
            | Def(_,_,_,_,exp) -> (extractfuncFromExpr exp extractArrayFromExpr joinMap)
            | _ -> Map.empty
    joinMap arrayMap t



let mutable defFunctions = Map.empty
let addDefFunction id var expr =
     defFunctions <- defFunctions.Add(id, (var, expr))
     Map.empty


let createFunctionTable (functionMap) stmt =
    let t = match stmt with
            | Def (_,_,id,var,exp) -> joinMap (extractfuncFromExpr exp extractFunctionFromExpr joinMap) (addDefFunction id var exp)
            | LetVar(_,_, _, exp) -> (extractfuncFromExpr exp extractFunctionFromExpr joinMap)
            | LetList(_,_,id, exp1, exp2) -> joinMap (extractfuncFromExpr exp1 extractFunctionFromExpr joinMap) (extractfuncFromExpr exp2 extractFunctionFromExpr joinMap)
            | LetTable(_,_,_,exp1, exp2, exp3) -> joinMap (extractfuncFromExpr exp1 extractFunctionFromExpr joinMap) (joinMap (extractfuncFromExpr exp2 extractFunctionFromExpr joinMap) (extractfuncFromExpr exp3 extractFunctionFromExpr joinMap))
            | If(_,_,exp1,_,exp2,_) -> joinMap (extractfuncFromExpr exp1 extractFunctionFromExpr joinMap)  (extractfuncFromExpr exp2 extractFunctionFromExpr joinMap)
            | Print(_,_, list) -> List.fold extractFunctionFromPrint Map.empty list
            | Read(_,_,list) -> List.fold extractFunctionFromRead Map.empty list
            | For(_,_,_,exp1,exp2, Some(exp3)) -> joinMap (extractfuncFromExpr exp1 extractFunctionFromExpr joinMap) (joinMap (extractfuncFromExpr exp2 extractFunctionFromExpr joinMap) (extractfuncFromExpr exp3 extractFunctionFromExpr joinMap))
            | For(_,_,_,exp1,exp2, None) -> joinMap (extractfuncFromExpr exp1 extractFunctionFromExpr joinMap)  (extractfuncFromExpr exp2 extractFunctionFromExpr joinMap)  
            | _ -> Map.empty
    joinMap functionMap t




let createMatrixTable (matrixMap) stmt =
    let t = match stmt with
            | Dim (_,_,list) -> List.fold extractMatrixFromDim Map.empty list
            | LetVar(_,_, _, exp) -> (extractfuncFromExpr exp extractMatrixFromExpr joinMap)
            | LetList(_,_,_, exp1, exp2) -> joinMap (extractfuncFromExpr exp1 extractMatrixFromExpr joinMap) (extractfuncFromExpr exp2 extractMatrixFromExpr joinMap)
            | LetTable(_,_,id,exp1, exp2, exp3) -> joinMap (joinMap (extractfuncFromExpr exp1 extractMatrixFromExpr joinMap) (joinMap (extractfuncFromExpr exp2 extractMatrixFromExpr joinMap) (extractfuncFromExpr exp3 extractMatrixFromExpr joinMap))) (Map.empty.Add(string id, (11,11)))
            | If(_,_,exp1,_,exp2,_) -> joinMap (extractfuncFromExpr exp1 extractMatrixFromExpr joinMap)  (extractfuncFromExpr exp2 extractMatrixFromExpr joinMap)
            | Print(_,_, list) -> List.fold extractMatrixFromPrint Map.empty list
            | Read(_,_,list) -> List.fold extractMatrixFromRead Map.empty list
            | For(_,_,_,exp1,exp2, Some(exp3)) -> joinMap (extractfuncFromExpr exp1 extractMatrixFromExpr joinMap) (joinMap (extractfuncFromExpr exp2 extractMatrixFromExpr joinMap) (extractfuncFromExpr exp3 extractMatrixFromExpr joinMap))
            | For(_,_,_,exp1,exp2, None) -> joinMap (extractfuncFromExpr exp1 extractMatrixFromExpr joinMap)  (extractfuncFromExpr exp2 extractMatrixFromExpr joinMap)
            | Def(_,_,_,_,exp) -> (extractfuncFromExpr exp extractMatrixFromExpr joinMap)
            | _ -> Map.empty
    joinMap matrixMap t

let checkLineNumberOrFail prevLineNum currLineNum line= 
   if currLineNum > prevLineNum then
      currLineNum
   else
      failwithf "Error. Label %d is not in incremental order; in line %d" currLineNum line

let lineNumList = new List<int>()
let dataList = new List<float>()
      
let labelExistsOrFail label lineNum =
    if not (lineNumList.Contains label) then
        failwithf "Error in line %d. Label %d does not exist." lineNum label

let checkExistingLabelsOrFail stmtList =
    for stmt in stmtList do
        match stmt with
        | If (num, _, _, _, _, label) -> labelExistsOrFail label num
        | Goto (num, _, label) -> labelExistsOrFail label num
        | GoSub (num, _, label) -> labelExistsOrFail label num
        | _ -> ()

let checkForNextOrFail stmtList =
    let stackWithLines = new Stack<(string*int)>()
    let stack = new Stack<string> ()
    for stmt in stmtList do
        match stmt with
        | For (num, _, id, _, _, _) ->
            if stack.Contains(id) then
                failwithf "Error in line %d. Variable %s used more than once in nested FOR" num id
            else
                stack.Push(id);
                stackWithLines.Push(id, num)
        | Next (num, _, id) ->         
            let ido,_ = stackWithLines.Pop()
            if stack.Pop() <> id then
                failwithf "Error in line %d. Bad FOR, expecting %s in NEXT" num ido
        | _ -> ()
    if stack.Count <> 0 then
        let mutable s = ""
        for id, line in stackWithLines do
            s <- (s + sprintf "Error. Missing NEXT %s statement for FOR in line %d\n" id line)
        failwithf "%s" s
 
let checkSemantic stmtList =
    let mutable numEnds = 0
    let mutable prevLineNumber = 0 
    let mutable numReads = 0
    
    for stmt in stmtList do
       match stmt with
       | End (line, lineNum) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           numEnds <- numEnds + 1; 
           lineNumList.Add(lineNum)
       | Read (line, lineNum, _) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           numReads <- numReads + 1;
           lineNumList.Add(lineNum)
       | Dim (line, lineNum,_) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | LetVar(line, lineNum, _, _) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | LetList(line,lineNum,_, _, _) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | LetTable(line, lineNum,_,_, _, _) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | If(line,lineNum,_,_,_,_) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Print(line,lineNum, _) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | For(line,lineNum,_,_,_,_) -> 
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Data(line,lineNum,list) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           list |> List.iter (fun data -> dataList.Add( toFloat data line));
           lineNumList.Add(lineNum);
       | Rem(line,lineNum) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Stop(line,lineNum) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Def(line,lineNum,_,_,_) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Goto(line,lineNum,_) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | GoSub(line,lineNum,_) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Return(line,lineNum) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
       | Next(line,lineNum,_) ->
           prevLineNumber <- checkLineNumberOrFail prevLineNumber lineNum line;
           lineNumList.Add(lineNum)
    
    if numEnds = 1 then
        match List.reduce (fun _ elem -> elem) stmtList with
        | End(_,_) -> ()
        | _ -> failwithf "Error. Last line is not an END statement"
    else if numEnds = 0 then
        failwithf "Error. Should have an END statement at the last line"
    else
        failwithf "Error. Should have only one END statement and it should be the last one"
    
    if numReads >= 1 then
        if dataList.Count = 0 then
            failwithf "Error. Should have at least one DATA available for READ"
    
    checkForNextOrFail stmtList
    checkExistingLabelsOrFail stmtList

let checkUsrFunOrFail k v =
   if (not (defFunctions.ContainsKey k)) then
       failwithf "Error in line %d. Function %s is not defined with DEF" v k 

    
let semanticAnalysis parseTree =
    match parseTree with 
        | Program (stmtList) -> 
            checkSemantic stmtList
            let functionTable = (List.fold createFunctionTable Map.empty stmtList)
            Map.iter checkUsrFunOrFail functionTable
            let scalarTable = List.fold createScalarTable Set.empty stmtList
            let arrayTable = joinMap (List.fold createArrayTable Map.empty stmtList) dimList
            let matrixTable = joinMap (List.fold createMatrixTable Map.empty stmtList) dimMatrix
            (scalarTable, arrayTable, matrixTable, defFunctions, dataList)
        
